import os
import sys
import gc
import random

import dgl
import dgl.function as fn

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch_sparse import SparseTensor
from torch_sparse import remove_diag, set_diag

import numpy as np
import scipy.sparse as sp
from sklearn.metrics import f1_score
from tqdm import tqdm

sys.path.append('../data')
from data.data_loader import data_loader

import warnings
warnings.filterwarnings("ignore", message="Setting attributes on ParameterList is not supported.")
warnings.filterwarnings("ignore", message="Setting attributes on ParameterDict is not supported.")


def set_random_seed(seed=0):
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed(seed)


def evaluator(gt, pred):
    gt = gt.cpu().squeeze()
    pred = pred.cpu().squeeze()
    return f1_score(gt, pred, average='micro'), f1_score(gt, pred, average='macro')


def get_n_params(model):
    pp = 0
    for p in list(model.parameters()):
        nn = 1
        for s in list(p.size()):
            nn = nn*s
        pp += nn
    return pp


# 这段代码定义了一个名为 hg_propagate_feat_dgl 的函数，用于在异构图（Heterogeneous Graph）中进行特征传播。该函数以 DGL 图为输入，逐跳地传播节点特征，最终返回更新后的图对象。
# g：DGL图对象（异构图）。
# tgt_type：目标节点类型，表示需要关注的节点类型。
# num_hops：传播的跳数，即特征传播的深度。
# max_length：元路径的最大长度，用于控制传播的范围。
# extra_metapath：额外的元路径列表，定义了特定的路径模式。
# echo：布尔值，决定是否打印调试信息。
# hg_propagate_feat_dgl 函数用于在异构图中逐跳地传播节点特征，并根据元路径和跳数控制特征的传播范围。
# 该函数的设计考虑到了元路径的保留条件，并对不同的节点类型分别进行处理，确保只保留需要的特征，从而优化内存和计算资源的使用。
def hg_propagate_feat_dgl(g, tgt_type, num_hops, max_length, extra_metapath, echo=False):
    # 循环遍历从 1 到 max_length-1 的跳数。每次循环处理一个跳数的特征传播。
    for hop in range(1, max_length):
        # 生成 reserve_heads，它包含了 extra_metapath 中那些长度大于当前 hop 值的路径的前 hop 个节点类型。
        # 就是让extra_metapath中每个元素的长度都不超过hop
        reserve_heads = [ele[:hop] for ele in extra_metapath if len(ele) > hop]
        # 遍历图 g 中的所有边类型 etype，并将其转换为标准的三元组格式 (stype, etype, dtype)，其中 stype 是源节点类型，dtype 是目标节点类型。
        for etype in g.etypes:
            stype, _, dtype = g.to_canonical_etype(etype)
            # if hop == args.num_hops and dtype != tgt_type: continue
            # 遍历源节点类型 stype 下的所有数据键 k，并检查其长度是否等于当前跳数 hop，以此判断该键是否对应于当前跳数的特征。
            # k就是g.nodes['M'].data['M'] = M中的'M'
            for k in list(g.nodes[stype].data.keys()):
                if len(k) == hop:
                    # 构造当前目标节点的特征名称 current_dst_name，它由目标节点类型 dtype 和键 k 组合而成。
                    current_dst_name = f'{dtype}{k}'
                    # 如果当前跳数为 num_hops 且目标节点类型 dtype 不是目标节点类型 tgt_type，并且该键不在 reserve_heads 中，则跳过该特征的传播。
                    # 同样，如果跳数超过 num_hops，并且键不在 reserve_heads 中，也跳过。
                    if (hop == num_hops and dtype != tgt_type and k not in reserve_heads) \
                      or (hop > num_hops and k not in reserve_heads):
                        continue
                    # 如果 echo 为 True，打印当前处理的键、边类型和目标节点名称。
                    if echo:
                        print(k, etype, current_dst_name)
                    # 使用 DGL 的 update_all 方法进行特征传播：
                    # 首先使用 fn.copy_u 复制源节点的特征 k 为消息 m，再使用 fn.mean 将消息聚合为目标节点的特征 current_dst_name。
                    # g 是一个 DGL 图对象。
                    # etype 是边类型（edge type）。
                    # fn.copy_u(k, 'm')：从源节点复制特征 k，并将其作为消息 m 发送到目标节点。
                    # fn.mean('m', current_dst_name)：将所有接收到的消息 m 在目标节点处取平均值，并存储到目标节点的新特征 current_dst_name 中。
                    # 1.消息传递 (fn.copy_u)：
                    # 对于边 etype，从源节点的特征 k 中提取信息，并将其作为消息 m 发送给目标节点。
                    # 2.消息聚合 (fn.mean)：
                    # 目标节点接收到来自所有相邻源节点的消息 m，然后将这些消息取平均值，并存储为新的特征 current_dst_name。
                    # 3.更新节点特征：
                    # 最终，目标节点的特征 current_dst_name 被更新为该节点从其所有邻居节点接收到的消息的平均值。
                    # todo k应该不是特征，而是元路径
                    g[etype].update_all(
                        fn.copy_u(k, 'm'),
                        fn.mean('m', current_dst_name), etype=etype)

        # remove no-use items
        # 经过一次跳数的特征传播后，清理不再需要的特征。遍历所有节点类型 ntypes，除了目标节点类型 tgt_type，将跳数小于等于当前 hop 的特征键 k 移除。
        # 如果 echo 为 True，则打印被移除的特征键列表。最后，调用 gc.collect() 进行垃圾回收。
        for ntype in g.ntypes:
            if ntype == tgt_type:
                continue
            removes = []
            for k in g.nodes[ntype].data.keys():
                if len(k) <= hop:
                    removes.append(k)
            for k in removes:
                g.nodes[ntype].data.pop(k)
            if echo and len(removes):
                print('remove', removes)
        gc.collect()

        # 如果 echo 为 True，则打印当前跳数的信息，包括节点类型、键、特征矩阵的形状，以及最后一列的最大值和均值。
        if echo:
            print(f'-- hop={hop} ---')
        for ntype in g.ntypes:
            for k, v in g.nodes[ntype].data.items():
                print(f'ntype:{ntype} k:{k} v.shape:{v.shape}', v[:,-1].max(), v[:,-1].mean())
        if echo:
            print(f'------\n')
    return g


# 这行代码定义了一个函数hg_propagate_sparse_pyg，用于在异构图上执行稀疏矩阵的标签传播。它接受多个参数：
# adjs: 存储不同类型边的邻接矩阵的字典。
# tgt_types: 目标节点的类型。
# num_hops: 传播的跳数。
# max_length: 元路径的最大长度。
# extra_metapath: 额外的元路径信息。
# prop_feats: 是否传播特征，默认为False。
# echo: 是否打印调试信息，默认为False。
# prop_device: 计算时使用的设备，默认为cpu。
# 这个函数主要通过矩阵乘法生成不同元路径对应的邻接矩阵，并根据设定的元路径长度限制和保留路径条件，构造出用于标签传播的特征矩阵。该过程涉及了大量的稀疏矩阵操作，并且利用了GPU加速（如果指定了GPU设备）。
def hg_propagate_sparse_pyg(adjs, tgt_types, num_hops, max_length, extra_metapath, prop_feats=False, echo=False, prop_device='cpu'):
    # 这行代码将存储设备设为cpu，即使计算在其他设备上（如GPU）进行，结果也会被移回cpu存储。
    store_device = 'cpu'
    # 这段代码检查tgt_types是否为列表，如果不是，则将其转换为包含单一元素的列表。这确保了后续代码可以统一处理tgt_types。
    if type(tgt_types) is not list:
        tgt_types = [tgt_types]

    # 这行代码创建一个label_feats字典，用于存储初始的标签特征。它从adjs字典中复制所有目标类型的邻接矩阵（即以目标类型结束的元路径），或根据prop_feats参数决定是否包含所有特征。
    label_feats = {k: v.clone() for k, v in adjs.items() if prop_feats or k[-1] in tgt_types} # metapath should start with target type in label propagation
    # 这行代码将adjs字典中的每个矩阵移动到指定的计算设备prop_device（通常是GPU或CPU），以提高计算效率。
    adjs_g = {k: v.to(prop_device) for k, v in adjs.items()}

    # 这行代码开始一个循环，遍历从2到max_length-1的每一跳。每一次循环都会在元路径上进一步传播特征。
    for hop in range(2, max_length):
        # 这行代码生成reserve_heads列表，其中包含长度大于当前跳数的额外元路径的后缀。这些路径将用于判断哪些新生成的元路径需要保留。
        reserve_heads = [ele[-(hop+1):] for ele in extra_metapath if len(ele) > hop]
        # 这行代码初始化一个空字典new_adjs，用于存储新生成的邻接矩阵。
        new_adjs = {}
        # 这行代码开始遍历label_feats中的每个项，即当前所有已生成的元路径及其对应的邻接矩阵。
        for rtype_r, adj_r in label_feats.items():
            # 这行代码将元路径的类型转换为列表，方便后续操作。
            metapath_types = list(rtype_r)
            # 这行代码检查当前元路径的长度是否等于当前跳数hop，如果是，表示这个元路径可以扩展。
            if len(metapath_types) == hop:
                # 这行代码提取当前元路径的起始类型dtype_r和结束类型stype_r。
                dtype_r, stype_r = metapath_types[0], metapath_types[-1]
                # 这行代码开始遍历所有的原始邻接矩阵，尝试将这些邻接矩阵与当前的元路径矩阵相乘以生成新的元路径。
                for rtype_l, adj_l in adjs_g.items():
                    # 这行代码提取当前邻接矩阵的起始类型dtype_l和结束类型stype_l。
                    dtype_l, stype_l = rtype_l
                    # 这行代码检查当前邻接矩阵的结束类型是否与元路径的起始类型匹配，如果匹配，表示可以相乘生成新的元路径。
                    if stype_l == dtype_r:
                        # 这行代码生成一个新的元路径的名称name，表示当前元路径通过该邻接矩阵扩展而成。
                        name = f'{dtype_l}{rtype_r}'
                        # 这段代码决定是否保留新生成的元路径：
                        # 如果当前跳数等于num_hops，并且新的元路径的起始类型不在目标类型中，且该路径不在保留路径中，则跳过。
                        # 如果当前跳数大于num_hops，且新的元路径不在保留路径中，则跳过。
                        if (hop == num_hops and dtype_l not in tgt_types and name not in reserve_heads) \
                          or (hop > num_hops and name not in reserve_heads):
                            continue
                        # 这段代码生成新的元路径矩阵：
                        # 如果name不在new_adjs中，则生成新的邻接矩阵。
                        # 如果在cpu上进行计算，直接相乘。
                        # 如果在gpu上，使用torch.no_grad()上下文避免计算图的构建，然后进行矩阵相乘，并将结果移回cpu。
                        if name not in new_adjs:
                            if echo:
                                print('Generating ...', name)
                            if prop_device == 'cpu':
                                new_adjs[name] = adj_l.matmul(adj_r)
                            else:
                                with torch.no_grad():
                                    new_adjs[name] = adj_l.matmul(adj_r.to(prop_device)).to(store_device)
                        else:
                            # 如果name已经在new_adjs中，则发出警告（如果echo为True）。
                            if echo:
                                print(f'Warning: {name} already exists')
        # 将新生成的邻接矩阵添加到label_feats中，用于下一跳的计算。
        label_feats.update(new_adjs)

        # 这段代码找出需要移除的元路径：
        # 如果元路径的起始类型在目标类型中，则保留。
        # 如果元路径的长度小于等于当前跳数，则标记为需要移除。
        removes = []
        for k in label_feats.keys():
            metapath_types = list(k)
            if metapath_types[0] in tgt_types: continue  # metapath should end with target type in label propagation
            if len(metapath_types) <= hop:
                removes.append(k)
        # 移除标记的元路径，减少内存占用。
        for k in removes:
            label_feats.pop(k)
        # 如果有移除的路径且echo为True，打印被移除的路径。
        if echo and len(removes):
            print('remove', removes)
        # 删除new_adjs字典，并手动调用垃圾回收器回收内存。
        del new_adjs
        gc.collect()

    # 如果计算设备不是cpu，删除adjs_g并清空GPU缓存，以释放显存。
    if prop_device != 'cpu':
        del adjs_g
        torch.cuda.empty_cache()

    # 返回生成的label_feats字典，其中包含了通过元路径传播后的特征。
    return label_feats


def check_acc(preds_dict, condition, init_labels, train_nid, val_nid, test_nid, show_test=True, loss_type='ce'):
    mask_train, mask_val, mask_test = [], [], []
    remove_label_keys = []
    k = list(preds_dict.keys())[0]
    v = preds_dict[k]
    if loss_type == 'ce':
        na, nb, nc = len(train_nid), len(val_nid), len(test_nid)
    elif loss_type == 'bce':
        na, nb, nc = len(train_nid) * v.size(1), len(val_nid) * v.size(1), len(test_nid) * v.size(1)

    for k, v in preds_dict.items():
        if loss_type == 'ce':
            pred = v.argmax(1)
        elif loss_type == 'bce':
            pred = (v > 0).int()

        a, b, c = pred[train_nid] == init_labels[train_nid], \
                  pred[val_nid] == init_labels[val_nid], \
                  pred[test_nid] == init_labels[test_nid]
        ra, rb, rc = a.sum() / na, b.sum() / nb, c.sum() / nc

        if loss_type == 'ce':
            vv = torch.log(v / (v.sum(1, keepdim=True) + 1e-6) + 1e-6)
            la, lb, lc = F.nll_loss(vv[train_nid], init_labels[train_nid]), \
                         F.nll_loss(vv[val_nid], init_labels[val_nid]), \
                         F.nll_loss(vv[test_nid], init_labels[test_nid])
        else:
            vv = (v / 2. + 0.5).clamp(1e-6, 1-1e-6)
            la, lb, lc = F.binary_cross_entropy(vv[train_nid], init_labels[train_nid].float()), \
                         F.binary_cross_entropy(vv[val_nid], init_labels[val_nid].float()), \
                         F.binary_cross_entropy(vv[test_nid], init_labels[test_nid].float())
        if condition(ra, rb, rc, k):
            mask_train.append(a)
            mask_val.append(b)
            mask_test.append(c)
        else:
            remove_label_keys.append(k)
        if show_test:
            print(k, ra, rb, rc, la, lb, lc, (ra/rb-1)*100, (ra/rc-1)*100, (1-la/lb)*100, (1-la/lc)*100)
        else:
            print(k, ra, rb, la, lb, (ra/rb-1)*100, (1-la/lb)*100)
    print(set(list(preds_dict.keys())) - set(remove_label_keys))

    print((torch.stack(mask_train, dim=0).sum(0) > 0).sum() / na)
    print((torch.stack(mask_val, dim=0).sum(0) > 0).sum() / nb)
    if show_test:
        print((torch.stack(mask_test, dim=0).sum(0) > 0).sum() / nc)


# model: 需要训练的模型。
# feats: 输入特征，可以是列表或字典格式。
# label_feats: 标签特征，也可以是字典格式。
# labels_cuda: 训练样本的标签，存储在GPU上。
# loss_fcn: 损失函数，用于计算模型预测与实际标签之间的误差。
# optimizer: 优化器，用于更新模型的参数。
# train_loader: 数据加载器，提供批量训练数据。
# evaluator: 评估器，用于计算模型的准确性。
# mask: （可选）掩码，用于指定哪些部分需要注意或处理。
# scalar: （可选）用于混合精度训练的标量缩放器。
def train(model, feats, label_feats, labels_cuda, loss_fcn, optimizer, train_loader, evaluator, mask=None, scalar=None):
    # 设置模型为训练模式：这行代码将模型设置为训练模式，使得在训练过程中启用如Dropout等特性。
    model.train()
    # 获取设备信息：获取labels_cuda所在的设备，通常是GPU设备。
    device = labels_cuda.device
    # 初始化变量：
    # total_loss: 用于累计每个批次的损失。
    # iter_num: 记录迭代次数。
    # y_true: 用于存储真实的标签。
    # y_pred: 用于存储模型的预测结果。
    total_loss = 0
    iter_num = 0
    y_true, y_pred = [], []

    # 开始训练循环：遍历训练数据加载器train_loader中的每个批次batch。
    for batch in train_loader:
        # 将批次数据移到设备上：这行代码被注释掉了，原本可能用于将批次数据移到设备上。因为后面的特征处理部分已经包含了这一操作，所以这里被注释掉。
        # batch = batch.to(device)
        # 处理输入特征：根据feats的类型，处理批次的输入特征：
        # 如果feats是列表，则为每个特征创建一个批次并移到设备上。
        # 如果feats是字典，则为每个键值对创建一个批次并移到设备上。
        # 如果feats既不是列表也不是字典，程序会抛出一个断言错误（assert 0）。
        if isinstance(feats, list):
            batch_feats = [x[batch].to(device) for x in feats]
        elif isinstance(feats, dict):
            batch_feats = {k: x[batch].to(device) for k, x in feats.items()}
        else:
            assert 0
        # 处理标签特征：将标签特征label_feats中的每个特征提取出当前批次的数据，并将其移到设备上。
        batch_labels_feats = {k: x[batch].to(device) for k, x in label_feats.items()}
        # 处理掩码：如果提供了mask，将其批次数据移到设备上；否则，batch_mask设为None。
        if mask is not None:
            batch_mask = {k: x[batch].to(device) for k, x in mask.items()}
        else:
            batch_mask = None
        # 提取当前批次的标签：从labels_cuda中提取当前批次的标签。
        batch_y = labels_cuda[batch]

        # 梯度清零：在每次反向传播之前，必须将优化器中的梯度清零，否则梯度会在不同批次之间累积。
        optimizer.zero_grad()
        # 混合精度训练：
        # 如果scalar不为None，表示要使用混合精度训练。
        # torch.cuda.amp.autocast()用于自动混合精度计算，这可以减少显存占用并加速训练。
        # 计算损失并进行反向传播，scalar.scale对损失进行缩放后再进行反向传播。
        # 更新优化器并缩放更新。

        # 条件判断：这行代码首先检查 scalar 是否为 None。scalar 是 torch.cuda.amp.GradScaler 的实例，用于在混合精度训练中对梯度进行缩放。
        # 如果 scalar 存在，则表示我们正在进行混合精度训练。
        if scalar is not None:
            # 自动混合精度上下文：torch.cuda.amp.autocast() 是 PyTorch 提供的上下文管理器，用于启用自动混合精度。
            # 它会自动将某些操作的浮点计算转换为半精度 (FP16)，而保持其他操作（如累加）使用单精度 (FP32)，从而在减少显存使用的同时保持数值稳定性。
            # 在 autocast 上下文中执行的操作，如模型的前向传递 (model(...))，会根据内部规则自动选择是使用 FP16 还是 FP32 进行计算。
            with torch.cuda.amp.autocast():
                # 模型前向传播：在 autocast 上下文中，调用模型的 forward 方法，计算模型的输出 output_att。由于 autocast 的存在，输入数据和计算过程中的部分操作将以半精度 (FP16) 进行。
                output_att = model(batch, batch_feats, batch_labels_feats, batch_mask)
                # 计算损失：仍然在 autocast 上下文中，计算模型输出与真实标签之间的损失 loss_train。autocast 会确保这个计算的数值稳定性，同时在可能的情况下使用半精度计算。
                loss_train = loss_fcn(output_att, batch_y)
            # 缩放损失并反向传播：scalar.scale(loss_train) 将损失 loss_train 进行缩放，目的是防止在反向传播过程中梯度值过小导致的数值不稳定问题。
            # 然后调用 backward() 方法，计算梯度。此时，梯度将根据缩放后的损失值进行计算。
            scalar.scale(loss_train).backward()
            # 应用梯度更新：scalar.step(optimizer) 进行优化器的参数更新。
            # 此操作将首先取消对梯度的缩放，然后再执行优化器的 step() 方法。这一步确保梯度的精度得到恢复，并且实际更新模型参数的梯度是未缩放的。
            scalar.step(optimizer)
            # 更新缩放因子：scalar.update() 更新 GradScaler 的缩放因子。这一步的目的是根据当前的梯度值动态调整缩放因子，以在后续的迭代中继续防止数值不稳定问题。
            scalar.update()
        else:
            # 常规训练：
            # 如果不使用混合精度训练，直接计算模型的输出output_att，并计算损失loss_train。
            # 进行反向传播并更新模型参数。
            output_att = model(batch, batch_feats, batch_labels_feats, batch_mask)
            loss_train = loss_fcn(output_att, batch_y)
            loss_train.backward()
            optimizer.step()

        # 记录预测结果和真实标签：
        # 将真实标签batch_y转换为长整型并移到CPU。
        # 根据损失函数的类型，选择相应的预测处理方法：
        # 如果使用二分类损失函数BCEWithLogitsLoss，将输出的概率大于0的部分视为预测的正类。
        # 否则，选择输出中最大值对应的索引作为预测标签。
        y_true.append(batch_y.cpu().to(torch.long))
        if isinstance(loss_fcn, nn.BCEWithLogitsLoss):
            y_pred.append((output_att.data.cpu() > 0.).int())
        else:
            y_pred.append(output_att.argmax(dim=-1, keepdim=True).cpu())
        # 累积损失并增加迭代次数。
        total_loss += loss_train.item()
        iter_num += 1
    # 计算平均损失：将累积的总损失除以迭代次数，得到平均损失。
    loss = total_loss / iter_num
    # 计算准确性：使用evaluator评估器计算真实标签和预测标签的准确性。
    acc = evaluator(torch.cat(y_true, dim=0), torch.cat(y_pred, dim=0))
    # 返回训练损失和准确性。
    return loss, acc


def load_dataset(args):
    # 使用 data_loader 函数加载指定路径的数据集。args.root 是数据集的根目录，args.dataset 是数据集名称。
    dl = data_loader(f'{args.root}/{args.dataset}')

    # use one-hot index vectors for nods with no attributes
    # === feats ===
    # 初始化 features_list 用于存储节点特征。
    # 初始化一个空的列表 features_list，用于存储节点特征。
    features_list = []
    # 遍历数据集中所有节点类型的数量。对于每种节点类型：
    # todo 不清楚dl.nodes['count']得到的是什么，待检查
    for i in range(len(dl.nodes['count'])):
        th = dl.nodes['attr'][i]
        # 如果节点没有特征（th is None），则使用单位矩阵表示其 one-hot 编码特征。
        if th is None:
            features_list.append(torch.eye(dl.nodes['count'][i]))
        # 否则，将节点特征转换为浮点型张量并添加到 features_list 中。
        else:
            features_list.append(torch.FloatTensor(th))

    # 初始化 idx_shift 数组，用于记录节点类型的起始索引位置。
    idx_shift = np.zeros(len(dl.nodes['count'])+1, dtype=np.int32)
    # 遍历所有节点类型，累加每种类型的节点数量，构建起始索引。
    for i in range(len(dl.nodes['count'])):
        idx_shift[i+1] = idx_shift[i] + dl.nodes['count'][i]

    # === labels ===
    # 获取训练标签中的类别数量。
    num_classes = dl.labels_train['num_classes']
    # 初始化 init_labels 数组，用于存储初始标签。维度为 (节点类型0的数量, 类别数量)（二维数组），类型为整数。
    init_labels = np.zeros((dl.nodes['count'][0], num_classes), dtype=int)

    # 获取训练/验证集和测试集的节点索引。使用 np.nonzero 函数获取标签掩码数组中非零元素的索引。
    # np.nonzero 是 NumPy 的一个函数，返回数组中非零元素的索引。
    # 对布尔数组应用 np.nonzero，会返回所有 True 元素的索引。
    # 对上面的示例数组 [False, True, False, True, True] 或 [0, 1, 0, 1, 1]，np.nonzero 返回 (array([1, 3, 4]),)，表示非零元素的索引为 1, 3, 4。
    # np.nonzero 返回的是一个包含索引的元组，元组中的第一个元素是索引数组。
    # [0] 表示取出这个索引数组。
    trainval_nid = np.nonzero(dl.labels_train['mask'])[0]
    test_nid = np.nonzero(dl.labels_test['mask'])[0]

    # 根据节点索引，将训练/验证集和测试集的标签数据填充到 init_labels 数组中。
    init_labels[trainval_nid] = dl.labels_train['data'][trainval_nid]
    init_labels[test_nid] = dl.labels_test['data'][test_nid]
    # 对于非 IMDB 数据集，将标签转换为类别索引（使用 argmax）。然后将标签数据转换为长整型张量。
    if args.dataset != 'IMDB':
        # init_labels.argmax(axis=1) 会沿着指定轴（这里是 axis=1，即每行）找到最大值的索引。
        # 这个操作会把 init_labels 转换为一个一维数组，长度为 num_nodes，其中每个元素是该行最大值的索引。
        # 换句话说，这一步将每个节点的多类别标签表示转换为单个类别标签（即每个节点被分配到一个特定的类别）。
        init_labels = init_labels.argmax(axis=1)
    init_labels = torch.LongTensor(init_labels)

    # === adjs ===
    # print(dl.nodes['attr'])
    # for k, v in dl.nodes['attr'].items():
    #     if v is None: print('none')
    #     else: print(v.shape)

    # 初始化 adjs 用于存储邻接矩阵。对于非 Freebase 数据集，使用列表存储；对于 Freebase 数据集，使用字典存储。
    adjs = [] if args.dataset != 'Freebase' else {}
    # 遍历数据集中所有的边类型和边数据。对于每种边类型：
    for i, (k, v) in enumerate(dl.links['data'].items()):
        # 将边数据转换为 COO 格式。
        # COO 格式是一种稀疏矩阵的存储格式，其中矩阵的非零元素由其行索引、列索引和对应的值构成。具体来说，COO 格式将稀疏矩阵存储为三个数组：
        # 行索引（row）：每个非零元素所在的行的索引。
        # 列索引（col）：每个非零元素所在的列的索引。
        # 值（data）：非零元素的实际值。
        v = v.tocoo()
        # 计算源节点和目标节点的类型索引。
        # v.col[0] 是一个数值，表示某个节点的全局索引。
        # idx_shift 是一个一维数组，通常表示节点类型的边界，长度为节点类型数加 1。
        # np.where 返回一个包含满足条件的元素索引的数组。
        # 这里的条件是 idx_shift > v.col[0]，意思是找到 idx_shift 中所有大于 v.col[0] 的元素的索引。
        # 假设 idx_shift 为 [0, 100, 200, 300]。
        # v.col[0] 为 150。
        # 则 np.where(idx_shift > 150) 返回 (array([2, 3]),)，因为 200 和 300 都大于 150。
        # [0] 提取第一个返回的数组。
        # [0] 再次提取该数组的第一个元素。
        # 这样就得到了 idx_shift 中第一个大于 v.col[0] 的元素的索引。（示例中得到2）
        # 最后一步，将索引减去 1。
        # 这一步的目的是得到源节点类型的索引，因为 idx_shift 通常是类型边界，而减 1 后的索引对应的是 v.col[0] 所属的类型。
        # todo 总感觉v.row[0]才是源节点的节点类型，待检查
        src_type_idx = np.where(idx_shift > v.col[0])[0][0] - 1
        dst_type_idx = np.where(idx_shift > v.row[0])[0][0] - 1
        # 根据节点类型索引，调整边的行列索引。
        row = v.row - idx_shift[dst_type_idx]
        col = v.col - idx_shift[src_type_idx]
        # 通过使用这些索引从 dl.nodes['count'] 中提取相应的节点数量，sparse_sizes 变成一个元组 (目标节点类型数量, 源节点类型数量)。
        sparse_sizes = (dl.nodes['count'][dst_type_idx], dl.nodes['count'][src_type_idx])
        # 创建稀疏张量 SparseTensor 存储邻接矩阵。其非零元素的位置由 (row, col) 指定。
        # SparseTensor 是一个用于创建稀疏张量的类。
        # row 和 col 分别是包含稀疏矩阵中非零元素的行索引和列索引的数组。它们被转换为 PyTorch 的 LongTensor 类型。
        # sparse_sizes 指定了稀疏张量的尺寸，即 (目标节点类型数量, 源节点类型数量)。
        adj = SparseTensor(row=torch.LongTensor(row), col=torch.LongTensor(col), sparse_sizes=sparse_sizes)
        # 对于 Freebase 数据集，将邻接矩阵存储在字典中；否则，存储在列表中。
        if args.dataset == 'Freebase':
            name = f'{dst_type_idx}{src_type_idx}'
            assert name not in adjs
            adjs[name] = adj
        else:
            adjs.append(adj)
            print(adj)

    # 对于 DBLP 数据集，构建异构图并设置节点特征数据。
    if args.dataset == 'DBLP':
        # A* --- P --- T
        #        |
        #        V
        # author: [4057, 334]
        # paper : [14328, 4231]
        # term  : [7723, 50]
        # venue(conference) : None
        # A, P, T, V 分别表示作者、论文、术语和会议的特征矩阵。
        # AP, PA, PT, PV, TP, VP 分别表示相应的邻接矩阵。
        # todo 这里为什么能直接赋值（对应上），待检查
        A, P, T, V = features_list
        AP, PA, PT, PV, TP, VP = adjs

        # 定义新的边集合 new_edges 和节点类型集合 ntypes。etypes 是定义的边类型，格式为 (源节点类型, 边类型, 目标节点类型)。
        new_edges = {}
        ntypes = set()
        etypes = [ # src->tgt
            ('P', 'P-A', 'A'),
            ('A', 'A-P', 'P'),
            ('T', 'T-P', 'P'),
            ('V', 'V-P', 'P'),
            ('P', 'P-T', 'T'),
            ('P', 'P-V', 'V'),
        ]
        # 遍历边类型和对应的邻接矩阵，提取邻接矩阵的行（dst）和列（src）索引，存入 new_edges，并将节点类型添加到 ntypes 中。最终使用 new_edges 创建 DGL 异质图 g。
        # zip(etypes, adjs) 将 etypes 和 adjs 打包在一起，形成一个元组列表，每个元组包含一个边类型和对应的邻接矩阵。
        for etype, adj in zip(etypes, adjs):
            # 将 etype 解包为 stype（源节点类型）、rtype（边类型）和 dtype（目标节点类型）。
            stype, rtype, dtype = etype
            # 将邻接矩阵 adj 转换为 COO（坐标）格式。adj.coo() 返回三个值：行索引（目标节点）、列索引（源节点）和数据（边的权重）。这里只使用了前两个值 dst 和 src。
            dst, src, _ = adj.coo()
            # 将 src 和 dst 从 PyTorch 张量转换为 NumPy 数组。这样做是为了便于后续操作。
            src = src.numpy()
            dst = dst.numpy()
            # 在 new_edges 字典中，以 (stype, rtype, dtype) 作为键，存储一个元组 (src, dst)。这个元组包含了源节点和目标节点的索引，表示从源节点到目标节点的边。
            new_edges[(stype, rtype, dtype)] = (src, dst)
            # 将源节点类型 stype 和目标节点类型 dtype 添加到节点类型集合 ntypes 中。ntypes 用于记录图中所有的节点类型。
            ntypes.add(stype)
            ntypes.add(dtype)
        # 使用 new_edges 构建一个 DGL 异构图 g。
        # new_edges 包含了所有边的信息，键是 (src_type, edge_type, dst_type)，值是对应的 (src, dst) 索引对。dgl.heterograph 会根据这些信息创建一个异构图。
        g = dgl.heterograph(new_edges)

        # for i, etype in enumerate(g.etypes):
        #     src, dst, eid = g._graph.edges(i)
        #     adj = SparseTensor(row=dst.long(), col=src.long())
        #     print(etype, adj)

        # g.ndata['feat']['A'] = A # not work
        # 将特征矩阵分配给对应类型的节点。
        # g.nodes 返回一个节点类型的访问器，可以通过节点类型名称（如 'A'）进行索引。
        # g.nodes['A'] 返回一个 NodeDataView 对象，这个对象允许访问和操作类型为 'A' 的节点数据。
        # .data 是 NodeDataView 对象的一个属性，用于存储节点的特征数据。
        # .data['A'] 是一个字典项，键是 'A'，表示该特征数据的名称。DGL 允许你为节点的特征数据命名，以便于后续访问和操作。
        g.nodes['A'].data['A'] = A
        g.nodes['P'].data['P'] = P
        g.nodes['T'].data['T'] = T
        g.nodes['V'].data['V'] = V
    elif args.dataset == 'IMDB':
        # A --- M* --- D
        #       |
        #       K
        # movie    : [4932, 3489]
        # director : [2393, 3341]
        # actor    : [6124, 3341]
        # keywords : None
        # 注释中描述了 IMDB 数据集的节点类型和边类型。
        # M, D, A, K 分别表示电影、导演、演员和关键词的特征矩阵。
        # MD, DM, MA, AM, MK, KM 分别表示相应的邻接矩阵。
        M, D, A, K = features_list
        MD, DM, MA, AM, MK, KM = adjs
        # 检查邻接矩阵的一致性。
        # 这句代码用于验证两个稀疏矩阵的列索引是否一致。具体来说，DM 和 MD 是两个稀疏矩阵，代码检查 DM 的列索引是否等于 MD 转置后的列索引。
        assert torch.all(DM.storage.col() == MD.t().storage.col())
        assert torch.all(AM.storage.col() == MA.t().storage.col())
        assert torch.all(KM.storage.col() == MK.t().storage.col())

        # 这句代码用于验证 MD 矩阵的每一行是否仅包含一个非零元素。具体来说，MD 是一个稀疏矩阵，代码检查每一行的非零元素个数是否为 1。
        # torch.all 是一个 PyTorch 函数，用于检查一个布尔张量中的所有元素是否都是 True。
        # 如果所有元素都是 True，则返回 True；否则返回 False。
        assert torch.all(MD.storage.rowcount() == 1) # each movie has single director

        # 同样地，提取邻接矩阵的行和列索引，存入 new_edges，并创建 DGL 异质图 g。
        new_edges = {}
        ntypes = set()
        etypes = [ # src->tgt
            ('D', 'D-M', 'M'),
            ('M', 'M-D', 'D'),
            ('A', 'A-M', 'M'),
            ('M', 'M-A', 'A'),
            ('K', 'K-M', 'M'),
            ('M', 'M-K', 'K'),
        ]
        for etype, adj in zip(etypes, adjs):
            stype, rtype, dtype = etype
            dst, src, _ = adj.coo()
            src = src.numpy()
            dst = dst.numpy()
            new_edges[(stype, rtype, dtype)] = (src, dst)
            ntypes.add(stype)
            ntypes.add(dtype)
        g = dgl.heterograph(new_edges)

        g.nodes['M'].data['M'] = M
        g.nodes['D'].data['D'] = D
        g.nodes['A'].data['A'] = A
        # todo 这里不知道是为什么
        if args.num_hops > 2 or args.two_layer:
            g.nodes['K'].data['K'] = K
    elif args.dataset == 'ACM':
        # A --- P* --- C
        #       |
        #       K
        # paper     : [3025, 1902]
        # author    : [5959, 1902]
        # conference: [56, 1902]
        # field     : None
        # 注释中描述了 ACM 数据集的节点类型和边类型。
        # P, A, C, K 分别表示论文、作者、会议和领域的特征矩阵。
        # PP, PP_r, PA, AP, PC, CP, PK, KP 分别表示相应的邻接矩阵。
        P, A, C, K = features_list
        PP, PP_r, PA, AP, PC, CP, PK, KP = adjs
        # 检查邻接矩阵的一致性。
        # P 是一个稀疏矩阵。
        # torch.where(P) 返回一个元组，包含两个张量 row 和 col，分别表示 P 中非零元素的行索引和列索引。
        row, col = torch.where(P)
        # PK 是另一个稀疏矩阵。
        # PK.storage.row() 获取 PK 矩阵中非零元素的行索引。
        # PK.storage.col() 获取 PK 矩阵中非零元素的列索引。
        # torch.all(row == PK.storage.row()) 检查 P 矩阵的行索引和 PK 矩阵的行索引是否完全一致。
        # torch.all(col == PK.storage.col()) 检查 P 矩阵的列索引和 PK 矩阵的列索引是否完全一致。
        # and 连接符表示两个条件必须同时为真。
        assert torch.all(row == PK.storage.row()) and torch.all(col == PK.storage.col())
        # AP 和 CP 是稀疏矩阵。
        # PK 是稀疏矩阵。
        # matmul 是矩阵乘法操作。
        # to_dense() 将稀疏矩阵转换为密集矩阵。
        # torch.all(... == A) 和 torch.all(... == C) 检查矩阵乘法的结果是否与给定的密集矩阵 A 和 C 一致。
        assert torch.all(AP.matmul(PK).to_dense() == A)
        assert torch.all(CP.matmul(PK).to_dense() == C)

        # PA, PC, PK, AP, CP, KP 均为稀疏矩阵。
        # .t() 是转置操作。
        # storage.col() 获取稀疏矩阵的列索引。
        # torch.all(... == ...) 检查转置矩阵和原矩阵的列索引是否一致。
        assert torch.all(PA.storage.col() == AP.t().storage.col())
        assert torch.all(PC.storage.col() == CP.t().storage.col())
        assert torch.all(PK.storage.col() == KP.t().storage.col())

        # 合并 PP 和 PP_r 邻接矩阵，并设置对角线元素。
        # PP 和 PP_r 是两个稀疏矩阵，分别表示某种关系（如 Paper-Paper）。
        # .coo() 方法将稀疏矩阵转换为坐标格式（COO），返回行索引、列索引和对应的值。
        # row0, col0, _ 获取 PP 矩阵的行索引、列索引和非零元素值（不使用值）。
        # row1, col1, _ 获取 PP_r 矩阵的行索引、列索引和非零元素值（不使用值）。
        row0, col0, _ = PP.coo()
        row1, col1, _ = PP_r.coo()
        # 使用 torch.cat((row0, row1)) 将 PP 和 PP_r 的行索引拼接起来。
        # 使用 torch.cat((col0, col1)) 将 PP 和 PP_r 的列索引拼接起来。
        # SparseTensor(row=..., col=..., sparse_sizes=PP.sparse_sizes()) 创建一个新的稀疏矩阵 PP，其中 row 和 col 分别是拼接后的行索引和列索引，sparse_sizes 保持与原 PP 相同的大小。
        PP = SparseTensor(row=torch.cat((row0, row1)), col=torch.cat((col0, col1)), sparse_sizes=PP.sparse_sizes())
        # PP.coalesce() 方法用于合并稀疏矩阵中的重复元素，并将其表示为规范的形式。合并后，每个位置上只有一个非零元素。
        PP = PP.coalesce()
        # PP.set_diag() 方法用于将稀疏矩阵的对角线元素设置为 1，确保每个节点都有自环（即对角线上的元素为1）。
        PP = PP.set_diag()
        # adjs 是一个包含多个稀疏矩阵的列表。
        # [PP] + adjs[2:] 将新的 PP 矩阵放在列表的第一个位置，并保留 adjs 列表中从第三个位置开始的所有元素。换句话说，adjs 列表的前两个元素被新的 PP 矩阵替换。
        adjs = [PP] + adjs[2:]

        # 同样地，提取邻接矩阵的行和列索引，存入 new_edges，并创建 DGL 异质图 g。
        new_edges = {}
        ntypes = set()
        etypes = [ # src->tgt
            ('P', 'P-P', 'P'),
            ('A', 'A-P', 'P'),
            ('P', 'P-A', 'A'),
            ('C', 'C-P', 'P'),
            ('P', 'P-C', 'C'),
        ]
        if args.ACM_keep_F:
            etypes += [
                ('K', 'K-P', 'P'),
                ('P', 'P-K', 'K'),
            ]
        for etype, adj in zip(etypes, adjs):
            stype, rtype, dtype = etype
            dst, src, _ = adj.coo()
            src = src.numpy()
            dst = dst.numpy()
            new_edges[(stype, rtype, dtype)] = (src, dst)
            ntypes.add(stype)
            ntypes.add(dtype)

        g = dgl.heterograph(new_edges)

        g.nodes['P'].data['P'] = P # [3025, 1902]
        g.nodes['A'].data['A'] = A # [5959, 1902]
        g.nodes['C'].data['C'] = C # [56, 1902]
        if args.ACM_keep_F:
            g.nodes['K'].data['K'] = K # [1902, 1902]
    # todo Freebase数据不知道为什么这样做，暂时先不管
    elif args.dataset == 'Freebase':
        # 0*: 40402  2/4/7 <-- 0 <-- 0/1/3/5/6
        #  1: 19427  all <-- 1
        #  2: 82351  4/6/7 <-- 2 <-- 0/1/2/3/5
        #  3: 1025   0/2/4/6/7 <-- 3 <-- 1/3/5
        #  4: 17641  4 <-- all
        #  5: 9368   0/2/3/4/6/7 <-- 5 <-- 1/5
        #  6: 2731   0/4 <-- 6 <-- 1/2/3/5/6/7
        #  7: 7153   4/6 <-- 7 <-- 0/1/2/3/5/7
        # 遍历 8 种节点类型，检查并打印邻接矩阵的信息。
        for i in range(8):
            kk = str(i)
            print(f'==={kk}===')
            for k, v in adjs.items():
                # k 是邻接矩阵的键，通常是一个元组 (t, s)，表示源类型和目标类型。
                # v 是对应的稀疏矩阵。
                t, s = k
                # assert s == t or f'{s}{t}' not in adjs 检查 s 和 t 是否相等，或者拼接 s 和 t 作为键是否不在 adjs 中。
                # 这是为了确保在字典中不会同时存在 (t, s) 和 (s, t) 两个键，确保图的表示一致性。
                assert s == t or f'{s}{t}' not in adjs
                if s == kk or t == kk:
                    # if s == t: 检查邻接矩阵是否对称，即源类型和目标类型相同
                    if s == t:
                        # 打印邻接矩阵的键 k，矩阵的大小（v.sizes()），以及非零元素的数量（v.nnz()）。
                        # 打印矩阵是否对称、对角线上的元素和总节点数的比率。
                        print(k, v.sizes(), v.nnz(),
                              f'symmetric {v.is_symmetric()}; selfloop-ratio: {v.get_diag().sum()}/{v.size(0)}')
                    # 如果源类型和目标类型不同（即矩阵不对称），仅打印 k，矩阵的大小和非零元素数量。
                    else:
                        print(k, v.sizes(), v.nnz())
        # 将 adjs['00'] 转换为对称矩阵。
        adjs['00'] = adjs['00'].to_symmetric()
        g = None
    else:
        assert 0

    # 根据数据集类型重新组织 adjs。
    if args.dataset == 'DBLP':
        adjs = {'AP': AP, 'PA': PA, 'PT': PT, 'PV': PV, 'TP': TP, 'VP': VP}
    elif args.dataset == 'ACM':
        adjs = {'PP': PP, 'PA': PA, 'AP': AP, 'PC': PC, 'CP': CP}
    elif args.dataset == 'IMDB':
        adjs = {'MD': MD, 'DM': DM, 'MA': MA, 'AM': AM, 'MK': MK, 'KM': KM}
    elif args.dataset == 'Freebase':
        # 创建一个新的字典 new_adjs。
        new_adjs = {}
        # 遍历 adjs 中的每个 (rtype, adj) 对：
        for rtype, adj in adjs.items():
            # rtype 是一个元组 (dtype, stype)，表示源类型和目标类型。
            dtype, stype = rtype
            # 如果源类型和目标类型不同（dtype != stype），则：
            if dtype != stype:
                # 创建一个新名称 new_name，拼接 stype 和 dtype。
                new_name = f'{stype}{dtype}'
                # 确保这个新名称 new_name 不在原有的 adjs 字典中。
                assert new_name not in adjs
                # 将转置矩阵 adj.t() 加入到 new_adjs 中。
                new_adjs[new_name] = adj.t()
        # 更新 adjs 字典，将 new_adjs 中的内容添加进去。
        adjs.update(new_adjs)
    # 如果 args.dataset 不匹配上述任何一个数据集名称，执行 assert 0，这会引发一个断言错误，表明数据集名称无效或未被处理。
    else:
        assert 0

    # 返回创建的异质图 g、邻接矩阵 adjs、初始化标签 init_labels、类别数量 num_classes、数据加载器 dl、训练和验证节点 ID trainval_nid、测试节点 ID test_nid。
    return g, adjs, init_labels, num_classes, dl, trainval_nid, test_nid


class EarlyStopping:
    def __init__(self, patience, verbose=False, delta=0, save_path='checkpoint.pt'):
        self.patience = patience
        self.verbose = verbose
        self.counter = 0
        self.best_score = None
        self.early_stop = False
        self.val_loss_min = np.Inf
        self.delta = delta
        self.save_path = save_path

    def __call__(self, val_loss, model):

        score = -val_loss

        if self.best_score is None:
            self.best_score = score
            self.save_checkpoint(val_loss, model)
        elif score < self.best_score - self.delta:
            self.counter += 1
            print(f'EarlyStopping counter: {self.counter} out of {self.patience}')
            if self.counter >= self.patience:
                self.early_stop = True
        else:
            self.best_score = score
            self.save_checkpoint(val_loss, model)
            self.counter = 0

    def save_checkpoint(self, val_loss, model):
        if self.verbose:
            print(f'Validation loss decreased ({self.val_loss_min:.6f} --> {val_loss:.6f}).  Saving model ...')
        torch.save(model.state_dict(), self.save_path)
        self.val_loss_min = val_loss
